import numpy as np


def build_causal_structure_from_abstract_state(pos_abs, pos_rela, vel_abs, vel_rela, action_dim):
    # mask的维度：（初始化为1）
    # (state-dim + reward-dim) * (state-dim + action-dim)
    state = pos_abs + pos_rela + vel_abs + vel_rela
    state_dim = len(state)
    action = [state_dim + i for i in range(action_dim)]

    mask = np.zeros([state_dim + 1, state_dim + action_dim])

    mask[0] = np.ones(state_dim + action_dim)

    for dim in pos_abs:
        mask[dim] = build_submask(state_dim, action_dim, [pos_abs, pos_rela, vel_rela, action])

    for dim in pos_rela:
        mask[dim] = build_submask(state_dim, action_dim, [pos_rela, vel_rela, action])

    for dim in vel_abs:
        mask[dim] = build_submask(state_dim, action_dim, [pos_abs, pos_rela, vel_abs, vel_rela, action])

    for dim in vel_rela:
        mask[dim] = build_submask(state_dim, action_dim, [pos_rela, vel_rela, action])

    return mask

def build_dense(pos_abs, pos_rela, vel_abs, vel_rela, action_dim):
    state = pos_abs + pos_rela + vel_abs + vel_rela
    state_dim = len(state)

    mask = np.ones([state_dim + 1, state_dim + action_dim])
    return mask

def build_submask(state_dim, action_dim, abstract_dim_list):
    submask = np.zeros(state_dim + action_dim)
    for abstract_dim in abstract_dim_list:
        for dim in abstract_dim:
            submask[dim] = 1

    return submask


def swap_rew_obs(mask):
    return np.roll(mask, shift=1, axis=0)
